{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "geek_chap_05.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pQx4HByM_IA6"
      },
      "source": [
        "##Sympy\n",
        "---\n",
        "Sympy is a python libraries prrmarily designed for symbolic calculation (and other things). Let us show that how we can use it to simply our calculation for logistic functions. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "t3Wjr0uW_Cry",
        "outputId": "1517764a-027c-46d6-93d9-8c131d40336e"
      },
      "source": [
        "!pip install sympy"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: sympy in /usr/local/lib/python3.6/dist-packages (1.1.1)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.6/dist-packages (from sympy) (1.1.0)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cs6VQ1GlAFzG"
      },
      "source": [
        "from sympy import *\n",
        "init_printing(use_unicode=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "egAgZB93At92"
      },
      "source": [
        "x = symbols('x')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3YVBww6SAIbB"
      },
      "source": [
        "def sigma(x):\n",
        "    return 1/(1 + exp(-x))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 44
        },
        "id": "J409ZWQnAqQf",
        "outputId": "cdaf81b4-4cf3-429f-f73e-6465183f7bcb"
      },
      "source": [
        "diff(sigma(x), x)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAEYAAAAkCAYAAAA0EkzVAAAABHNCSVQICAgIfAhkiAAAAtlJREFUaIHt2V+IFlUYx/HPamJmIGqQtWTK4oWgtbJZEe2ySMXi9XopeO+N6VV7EV0l0V0QGiVteVXRH7rboCVFS8LVVVQQBK9WE/+luxS00HbxzMI4vTvO67vzzlrzhcM753nnOeeZZ86ZM+c31DRkcdUBFOQ4zmMCh9CFE2V2+LAk5gr2oBNPYKjsDjvK7qAg43ikgf0NkRQYw+8YwHSb4iqdTnyOm+LivsaTTfhvxWV8M/+hVcd6XMN+bEQ3juDbgv6dOJO0cwqbSoixEkbwbsb2Gu4W8F2Gn9GX1Hfgi/kLrTqexQz+wFSq/Cmm1YKl0QNvPnlejIyeBv/9VXLfLVF2YqaxHL+JkVKTsBLXxYN2i3gxex0fYlGFcS0IXsCoWKYnxQrzVqUR1dTU1FTM7CZyptIoahYMz+AnXMBZsdWowVNiMwtrhAC2PM9hpdgJd5UbV2kcx0vJ8SG8mRx/hX05fmfEKJqT9/Fpqt6H70VGZ7Cr+VjbyoB4y96Lj1P2zbiFFQ18enAur9HHcBuvpmzbhWQwKHbIu5oMdBjvNOkzF+PiArLl6cx5Y/gRSzL2k9idsa0SWvIreR0PiqzOJXdOKTcxrap85Ct9b+NYqr4UR7GzUUPpjVyvyHYVS/esOjchRmy/EL0PNtFGJz7BNqzzb6XvV7woxK8OcdNGcfh+DX+Hz3L+L3PEtKLyUUzpe07c9C6R/L/F9Jwtm/OC+yin8yKJGXKvUjctBKm0rTfj0y6Vb0PSTyHNOC1U3RDLdSscxJep+ntienyQsk1kfNql8q1Kfq8XOTmdmNNaX45vJWWWyaR+KcenXSrfJnFTrhU5Of3wHRGfN1anbI+LN8Tu5Ny1yfHa+Yg04YR4TTisXJWvV1zjA/GLe9f6fjEvs2W4YHvDij18y1b5HsUdvPygDQzgoofnm3ZRduOHZhyyCbgkhu9VkeH/CltwwAL/llVTU/P/4R8zrbCaKizkdAAAAABJRU5ErkJggg==\n",
            "text/latex": "$$\\frac{e^{- x}}{\\left(1 + e^{- x}\\right)^{2}}$$",
            "text/plain": [
              "    -x    \n",
              "   ℯ      \n",
              "──────────\n",
              "         2\n",
              "⎛     -x⎞ \n",
              "⎝1 + ℯ  ⎠ "
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LAyAnSWaAsUp"
      },
      "source": [
        "def log_lik(t):\n",
        "    return log(sigma(t))\n",
        "def log_lik_minus(t):\n",
        "    return log(1-sigma(t))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 41
        },
        "id": "loDK1YWNBHj4",
        "outputId": "9b0fb16a-8219-4998-b2fe-48b32b53aef8"
      },
      "source": [
        "diff(log_lik(x),x)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAADQAAAAiCAYAAAAQ9/ptAAAABHNCSVQICAgIfAhkiAAAAdxJREFUWIXt2E+ojFEYx/GPy0ZslFKm/MlKKSPsuN2EblnbW9uQnbuyIlulbqEmVhTK6tooC7obuoqdshqR/Mm9UWYxFuedeo2XOe/7zjvvpfdbp+mcM8/p93TOPOf8hv+MtXULwBO8Qhc3sAuLRRdbDQm9xVm0sBlzZRZbMw5FI1jCuozx40Iy8AxfMIveBDT9lRZu4qMg6i625Ig/iDe4N35p+dmJ97iE3WjjMe5HxrfwIlnnOfZUoDEXD3FxaOwovkbErsdTTCf9k7g9Pmn52Y4+vmEl1b4Lx68Wsn6ssewVdmJ/xtyPEuuWokxCPWzAO2Fn/nk24YNQAPYJF+IxXMVUjbpKcQCPhHK9LFSs87UqamhoKMXgcdqvVUXDn1kNfqgIUaZwGg+SL/VxanL6cjMrXOjncC09kb7RN+IlzggPzDx0cKGwvN9ZSrQMt63J/AK24QROxyy4It8OdeRLqDJTWMeba2DmujiEGeG/hPnI+Bau4wh2iDSFVe5QLaZwVEJzfjV1PcEDpccOZ8RVbgqL+qF53En1LwtH6EpqrJsRV7kpLJrQp6QNWE76r0fEVW4Kh8t2O2lTQllsJ5/jYhGfccsETOGMcL6HWycitiO+KDSmsKGhYXL8BG/OgifFgx+bAAAAAElFTkSuQmCC\n",
            "text/latex": "$$\\frac{e^{- x}}{1 + e^{- x}}$$",
            "text/plain": [
              "   -x  \n",
              "  ℯ    \n",
              "───────\n",
              "     -x\n",
              "1 + ℯ  "
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49
        },
        "id": "pxFDrMLWBUUT",
        "outputId": "e05aeba6-397e-4abf-e242-3c54201914ac"
      },
      "source": [
        "diff(log_lik_minus(x), x)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAKMAAAAqCAYAAADf0+uQAAAABHNCSVQICAgIfAhkiAAABP9JREFUeJzt3GmoVVUUwPHfM3PIRDJpsl6FCQXagI2QdYkGbfpUSFQmFAVFNH2IioYPklRfbKRJemUEjYZWZIiUZFk0z0EzPUuyWSp6oH1Y99J9x+u797177j1v2H/Y8M4+e1juu886e62zliQSg4TtihYgYS0+RjeWYBrWFSpRQaTNWDzrcTmmYgquLVac4ugoWoBhznsYXaP+RLEJK7yN3zAHPW2QKzEMmYpH8LPYTE9j136OcRi+xjP5ipYYSeyLDViEA3AwXsGyfowxFe+Xx3oHM3KWMTFCWImbM3XH448G+4/HazimfH0mHs9HtMRIYm9swV/YVFX+Fq/sxACodbhO1OcgoQFn1bj3b5tlGTakzTgwejABPwqNmEgUxk74SRgrhwhH9Qm4G6MKlCsxQjkUq4VL509hFV9TqESJRKI1fCOsxEbLo4VImRiWZA2YL/FPP/qvr98kkUgkhhh5BUpsyWmcRGLEswy/4qmiBRnJJJ9YcDvmFy3EMGcvvIxP8IH4Ft8WJuLodk2WEyWt0YxzWzDmUGR3EdkEu4nI9gnVDVqhGTtwPd5owdhDkW6cv417O4kwtGntEydX1uKI8t9LcEXVvSdxVdX1DyLYmPiMuhGTWy3gfJyeqTsGy8UPswULWi3EAChp3ZnxLrWDbm/DQ5m6obBWFeaI8/aVeCBzbyZ+waQa/Wbho2xl3ppxDC7BC5n6HcuTXybCrEYaT+DGTN0OuEBolGqaXasu3DSAfrV4ryxLtuxRvv8iOnEKLs70/RBf4ZxM/WQRHX9hTjJuk7lYWqfNJoPzaS9pnWYcJTTd9lV1ZwjN0Zd7bSBr1aXxzdhs2kS9lIkb8GrV9Viswbm1GuetGU8V4fNDjVXijHMyvsdROY+/Gd9idlXdbJGIVZSPtpLq0C2MzZLITry3wf5T8SCOwz5qp0y8icNFVHuHeFBWq6+wcuEtobL7YrBqxlbTpXdUz7N4uE6fVmrGZtImGk2ZOFA8bNPEht8sXv2VMrOBuQZMt95Pfy0aWeCF6gdplJqQswjuwOKq65W4r06fRtbqWr1TH3pEtHl1XfY3aVfaxPTyPA0lmuUd6T1FxPY1y2L1I4K+y2GedvKHeDVW2ChcO81yrzCQKtwilMIdVXXdmT7tSpuouG5+aqRx3ptxkzikNsvGcukPg/H7eLVxMk7vB/Vd+RxXfimXCn+Wr7/oo0+70iZmiAdhQyON8zZgNqjtV9pReN8PLs/ZWf67M8e5O8Q5rGIVdrSx7I77sZ949W5nayt5kt4/ykqRb71zpl071mqd+Ba/VGvTJmaLf2chLMd5NepLap/7unKev4TT1HfRLDAwrXQQnsuUirtmkThajNtG3xU4O1P3uvDLVlPS3Fp1acyAaXXaxDj8jiNzHLNfXKT3Ib0ISprfjNPxvHC9rMEudcabKDZiVx9tvrW1FpyDzw3P/4DrErzUnw55v6ZXCL/SUGYs7hEP1iw8pu+vBaNF1M81Im2jllHQKb5GZC3VF8Wrcc+mJB6c9ODSooW4E/sXOH9Jbc04xv/+re/KpXI9pqrdPJFOUbn3Ka5uUqbrcGyTYyQGwBRxfiqKkuZe0wtxVn7imCA0Z6IOrQgh2ygMmSPqNRyk/IiTqq4PbHK8eSKkLlEg4wuYc5VwsP6l72/MC2xbM04QLqLPhIV5a5MyFbEOiUQikUgkEolEIpFoJf8BHMcz+XHAQ2QAAAAASUVORK5CYII=\n",
            "text/latex": "$$- \\frac{e^{- x}}{\\left(1 - \\frac{1}{1 + e^{- x}}\\right) \\left(1 + e^{- x}\\right)^{2}}$$",
            "text/plain": [
              "           -x           \n",
              "         -ℯ             \n",
              "────────────────────────\n",
              "                       2\n",
              "⎛       1   ⎞ ⎛     -x⎞ \n",
              "⎜1 - ───────⎟⋅⎝1 + ℯ  ⎠ \n",
              "⎜         -x⎟           \n",
              "⎝    1 + ℯ  ⎠           "
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f0AOwEu0Bk7Y"
      },
      "source": [
        "The results above are ugly. Let us try to simplify it using `simplify` function."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 39
        },
        "id": "MeQnfBtqBbSO",
        "outputId": "f30a0fba-2d11-4a09-c0c7-29f07e8005d3"
      },
      "source": [
        "simplify(diff(log_lik_minus(x), x))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAD0AAAAiCAYAAADs4tGnAAAABHNCSVQICAgIfAhkiAAAAfFJREFUaIHt2D9oFEEUx/GP/woxjSJYHFwUK0EwonYagkQQrO2tbRQ7U1kptoIQUOFQG0URUogWChZKGgVBOwWbC4r4N0HFA89iNrIed7rhZm4Tbr8wsDsz+5v39s/beY8hZE3ZBhTgMV6iiSvYjtl+BFeC03M4iRo2Y6pccwbHUzzAuhhiq2OI9EkNV/EBn3EbW3Lj+7AJX9AauHUJ2IZ3OIcdGMMj3MnGa3iezXuGnSXYGJ37ONvRN4mvWI8nGM/6j+LG4ExLwyja+IaFXPsuvOrJWJtS/D/sEp7oni5jP1MuXKbTLWzAW+EJDwUb8V4IWruFTcchXLQ8/irJ2IuHwq9qXojUp0u1aFh4I0TUou16KVb2SWcge40fS7h+LqItFSlZFUmnHUmnIhWDzqejFwRWAoeFzcgpXEqgP44Z4aa2cazbpEHvfO6hjiM4XvCaBs4UnDuCFzghJC5die102QWBu0I56RZ+9ZoU0+nFRL+J/ZgQalrT2XgNl3EQW1UFgX/SUPz1zrOgxzcdi5gFgakOjZaQX+f7DhTQ6el0rHw6ZkFgGjdz5+eFT+ZCrq+5RM2/iOV0zILAx6wtMp+dv+pT9w+xAtksPuGacgsCI0JFdSxbt54d11MtmKog0FA8kE3ongI3IthRUVFRsXz4DefRiiu+S5ywAAAAAElFTkSuQmCC\n",
            "text/latex": "$$- \\frac{e^{x}}{e^{x} + 1}$$",
            "text/plain": [
              "   x  \n",
              " -ℯ   \n",
              "──────\n",
              " x    \n",
              "ℯ  + 1"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pfv_81-LEXJP"
      },
      "source": [
        "Let us use sympy to simplify certain notations. Pay attention that we have used $m$ to replace the $x_i^t\\beta$ part in the lecture notes. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0Do4ZT0mBkDG"
      },
      "source": [
        " m, y = symbols('m y')\n",
        " def my_expression(m, y):\n",
        "     return y * (exp(-m)/(1+exp(-m))) - (1-y)*(exp(m)/(1+exp(m)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 41
        },
        "id": "tWLsfdvHFctM",
        "outputId": "f04346ef-93e4-4228-b3cc-f77278810029"
      },
      "source": [
        "simplify(my_expression(m, y))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAF0AAAAkCAYAAADrXDbKAAAABHNCSVQICAgIfAhkiAAAA2xJREFUaIHt2UuIHEUcx/FP1oivBDEKgo8EDSKaXaIYRcSVARVEELwqQgTFgwEVPUhyiaAI4iUGlBUUB/Xk42AQYRUED5KoxPjCIATFR6LBGB8bEk1EPVSN1rY9M90z3TN59BeK7qr6V/2r/1NdVdM/Go4JTsNuLB/3QHJ4GQ8M2HYTnsAWfIEr8Bq+xt2VjG4IHsdzY/B7jRCYnfgbt+fYTGEvTh2g/69wR7x/Cu9gMc7HB6nhxACd59HGQwXsTsadeLYiv2VYhM9wLw50sfkUX+K2kn0vxgLzn2sj5mL5b6lxGvTvcH+msyn8jotLDqIbNwqz7N1M+Sh8v4F1eAV/9bDbhFtyys/G8/gJv+BVnBnrVpg/m6fwXryfFH7Mf0mDvhmXZxxtwDP4vMcgyzCNrULgU0bhuyjvC+vxSUnZefhQWJquRgtnYCbWT+GTxP4cYSJ16goH/WZcivVDPECWZdiVUz4K30XZheNxVlI2Iywda7EdH+FhXBvr06Cfi2+Ttv8LespVwgxcghOwA/d1sV2HfUk6hIOZsumcdrN4ekjfj0TbXqnV7SEj++RvpHBB7GMy5pfF/H7zn++AsNSUZmFyv1UI3Cphlv2JJ7u0m8FLSf4x4dXbmJTtzGm3RzgyZinjewNe7FLX4Zs+9b1YEq8/xutKYSO8LMf24CAO0qD/gW24CatxqzCD89gbU4e5mN/Rx982+TOsjO89MdXFpDBhdsf8IZyCH4QZPjTZI+NmrInX16twkGEWF+H0nLq6fS/CJTFNYGm8X5qxm47j7LAFP+MF4S1cjuuFN7GSI/dq4dVeUbJdW7FzOv8FtyrfRWnJX//bic2J+BVXZtquwtvCUXEOHwubaiW8pftaWhU3CH+TjxuD736swZujcDQhHPIfxPfyN7qquUc4FYzDdy/uwoWjcNQS/qFtF45uo2ScvhsaGhoaGoZmQbxmv/o1HAMc1lLX0UphqetooCq5bhhKSV0VUEQrrZXDIeilpK4etBX7/lNEK62VUQW9l75YSuqqgKJaaW2MIuhl9MVSUldDd2bxaKbsOtWv1W3FPy936CXbHbFUri8mDKrTpowl6Av7mwxF5fpiwqA67dipO+iV64sJg+q0Y6fujbR2fXEAimqlRzS16osJbcU20pb+WmlDQ0NDQ0Mh/gFfYxFod98DZAAAAABJRU5ErkJggg==\n",
            "text/latex": "$$\\frac{y + \\left(y - 1\\right) e^{m}}{e^{m} + 1}$$",
            "text/plain": [
              "             m\n",
              "y + (y - 1)⋅ℯ \n",
              "──────────────\n",
              "     m        \n",
              "    ℯ  + 1    "
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OIBaYZMk-GSW"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uusXyKBU3qgL"
      },
      "source": [
        "## Jax time\n",
        "\n",
        "Now let us try to use Jax to derive a useful version of optimize routine. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1PfZvTfyFjGG"
      },
      "source": [
        "import jax \n",
        "from jax import lax\n",
        "from jax import random\n",
        "import jax.numpy as jnp"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tqxa49WrnZjd"
      },
      "source": [
        "# Question: Why define it like this\n",
        "def sigmoid(x):\n",
        "    return 0.5 * (jnp.tanh(x / 2) + 1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gXsKfkvv5pmg",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "8b444f49-0e93-487e-a766-a09e2175c7a2"
      },
      "source": [
        "key = random.PRNGKey(0)\n",
        "x_key, beta_key, beta_test_key = random.split(key,3)\n",
        "x = random.normal(x_key, (10000, 10))\n",
        "beta = random.normal(beta_key, (10,))*2.0\n",
        "beta_test = random.normal(beta_test_key, (10,))\n",
        "y = (sigmoid(x.dot(beta))>=0.5).astype(jnp.float32)\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "owM7l8wHnPBe",
        "outputId": "ba8b4dd4-ee3b-439a-c168-0b777ba45b00"
      },
      "source": [
        "y.sum() # Check whether the results are reasonable"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(5017., dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rRVc0iQmm83D"
      },
      "source": [
        "def predict(beta, x):\n",
        "    return sigmoid(x.dot(beta))\n",
        "\n",
        "def loss(beta):\n",
        "    preds = predict(beta,x)\n",
        "    label_probs = preds * y + (1 - preds) * (1 - y)\n",
        "    return -jnp.sum(jnp.log(label_probs))/10000.00 "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DWGZCgYkF_I2",
        "outputId": "1b396a7e-70f2-4a98-c313-61f8a06ef403"
      },
      "source": [
        "loss(beta)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(0.11576849, dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UwCE_H0noOUi",
        "outputId": "474b5210-26a6-40c2-90d9-fbbfea7543a6"
      },
      "source": [
        "grad_func = jax.grad(loss)\n",
        "grad_func(beta)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([ 0.00455573,  0.01565121, -0.00860528, -0.00485264,\n",
              "              0.00207953, -0.00256935,  0.00626604, -0.00537516,\n",
              "              0.0021444 , -0.00108901], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SdjRZxpkp5JZ",
        "outputId": "659ebba8-1fcc-4242-975d-2eaaef58010c"
      },
      "source": [
        "jax.grad(loss)(beta_test)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([ 0.24341547,  0.29687792, -0.02719581, -0.13187408,\n",
              "             -0.05313727, -0.11881256,  0.1569478 ,  0.10492601,\n",
              "             -0.05401436, -0.03668952], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpjKSDRJpL2K"
      },
      "source": [
        "# Question: What can We See from This Implementation?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3i6jWOkhpVOH"
      },
      "source": [
        "The derivatives should be close to $0$. Because we are evaluating it at the true function. It should never be equal to $0$ because we generally never have unlimited sample for that, however, it should not be too huge as well. \n",
        "\n",
        "Also, it should in general be different from, say, a randomly initialized beta. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "37MrrD_RptKk"
      },
      "source": [
        "Now let us try to make our own implementation of the grad function. It should be close to what autograd gives us. \n",
        "\n",
        "What we have derived is that the gradient is \\begin{align*}\n",
        "-\\sum_i (y_i - \\sigma(x_i^t\\beta) )x_i \n",
        "\\end{align*}\n",
        "\n",
        "This isn't the most convenint form for implementation, thanks to the pesky summation. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tdAtFsqUrTP3"
      },
      "source": [
        "# Question: How Can we Simplify the Above Expression?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KiS2vSGwpFMK"
      },
      "source": [
        "# Why this makes sense?\n",
        "def custom_grad(beta):\n",
        "    residual = y - predict(beta, x)\n",
        "    return jnp.transpose(x).dot(-residual)/10000.00"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TH0fjd0PqYIx",
        "outputId": "8dc45f69-581a-436d-9569-75401353fe4a"
      },
      "source": [
        "custom_grad(beta)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([ 0.00455572,  0.01565119, -0.00860529, -0.00485265,\n",
              "              0.00207953, -0.00256935,  0.00626603, -0.00537516,\n",
              "              0.00214439, -0.00108901], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AtNGTvS0qZ3H",
        "outputId": "6481e018-8f6b-4998-faaa-18e6c216b692"
      },
      "source": [
        "custom_grad(beta_test)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([ 0.24341445,  0.29687792, -0.02719574, -0.13187398,\n",
              "             -0.0531371 , -0.1188125 ,  0.15694797,  0.10492576,\n",
              "             -0.05401447, -0.03668943], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TyOwatC22Fcd"
      },
      "source": [
        "## Proximal methods implementation.\n",
        "---\n",
        "\n",
        "Here we implement the proximal methods and see whether it will lead to the desired outcome. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "roiHZPXYrdvL"
      },
      "source": [
        "def soft_threshold(x, thres):\n",
        "    return jax.lax.cond(x > thres,\n",
        "                        lambda _: x - thres,\n",
        "                        lambda _: jax.lax.cond(\n",
        "                            x < -thres,\n",
        "                            lambda _: x + thres,\n",
        "                            lambda _:0.0,\n",
        "                            None\n",
        "                        ),\n",
        "                        None)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SPND8vL23fmF"
      },
      "source": [
        "def proximal_methods(beta_init, max_iter, eps, lr, penalty):\n",
        "    converged = False\n",
        "    beta_old = beta_init\n",
        "    beta_new = beta_init\n",
        "    soft_threshold_partial = lambda x: soft_threshold(x, lr*penalty)\n",
        "    current_iter = 0\n",
        "    while not converged and current_iter < max_iter:\n",
        "        print(\"Current iteration is %d\"% current_iter)\n",
        "        beta_copy = beta_old \n",
        "        current_loss = loss(beta_copy) + penalty*jnp.linalg.norm(beta_copy, 1)\n",
        "        current_grad = custom_grad(beta_copy)\n",
        "        w = beta_copy - lr*current_grad\n",
        "        beta_new = jax.vmap(soft_threshold_partial, 0)(w)\n",
        "        new_loss = loss(beta_new) + penalty*jnp.linalg.norm(beta_new, 1)\n",
        "        diff = jnp.abs(new_loss-current_loss)\n",
        "        print(\"The difference is %.5f\"%diff)\n",
        "        beta_old = beta_new\n",
        "        if diff <= eps:   \n",
        "            converged = True\n",
        "            print(\"Algorithm converged\")\n",
        "            break\n",
        "        else:\n",
        "            current_iter +=1\n",
        "            if current_iter >= max_iter:\n",
        "                print(\"The algorithm have failed to converge.\")\n",
        "                break\n",
        "\n",
        "    return beta_new, converged"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L-7ATXj93xtU"
      },
      "source": [
        "key = random.PRNGKey(0)\n",
        "x_key, beta_init_key = random.split(key,2)\n",
        "x = random.normal(x_key, (10000, 5))\n",
        "beta = jnp.array([2.0,2.0,0.0,0.0,0.0])\n",
        "beta_init = random.normal(beta_init_key, (5,))\n",
        "y = (sigmoid(x.dot(beta))>=0.5).astype(jnp.float32)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PaXJfUhN4cac"
      },
      "source": [
        "Before we run the whole function, it is better to evaluate using a step-by-step approach to see where it might go wrong."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6J5mji1r4p6k"
      },
      "source": [
        "lr = 0.1\n",
        "penalty = 1.0\n",
        "eps = 1e-5"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tz6dlXiO322q"
      },
      "source": [
        "converged = False\n",
        "beta_old = beta_init\n",
        "beta_new = beta_init\n",
        "soft_threshold_partial = lambda x: soft_threshold(x, lr*penalty)\n",
        "current_iter = 0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eXFRD_Mt4VyX"
      },
      "source": [
        "beta_copy = beta_old \n",
        "current_loss = loss(beta_copy) \n",
        "current_grad = custom_grad(beta_copy)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oowhL3wD6SF8",
        "outputId": "7466f406-153d-4abb-95e7-eaa8e0cba13e"
      },
      "source": [
        "current_loss + penalty * jnp.linalg.norm(beta_copy,1)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(10.273039, dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 82
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sPJ4ECEo5mat",
        "outputId": "f397d321-e973-48e5-bcf0-f55a42483263"
      },
      "source": [
        "beta_copy"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([-1.4581941, -2.0470448,  2.0473385,  1.1684095, -0.9758365],            dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 83
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wxkoHWSm4m62"
      },
      "source": [
        "w = beta_copy - lr*current_grad"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mwwe6weX6CZs",
        "outputId": "ecfab29e-4af2-4455-88c3-d827ec0485ac"
      },
      "source": [
        "w"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([-1.4146639 , -1.9989496 ,  2.0274482 ,  1.1572018 ,\n",
              "             -0.96532995], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 85
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "350EIzSp4pUd"
      },
      "source": [
        "beta_new = jax.vmap(soft_threshold_partial, 0)(w)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "I2zL1KIm5sVC",
        "outputId": "d142eae7-f7aa-4c7f-c2f9-710e8384723b"
      },
      "source": [
        "beta_new"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([-1.3146639, -1.8989496,  1.9274482,  1.0572017, -0.8653299],            dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 87
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5k4RK2od5ywV"
      },
      "source": [
        "new_loss = loss(beta_new) + penalty*jnp.linalg.norm(beta_new, 1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3cH3Sre47sIz",
        "outputId": "769e6940-5e30-4d68-c270-38c48ba4ec55"
      },
      "source": [
        "new_loss"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(9.459178, dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 89
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l1YEzbCP7vgT"
      },
      "source": [
        "The above results seem reasonable, now let us test what happens near the optimal beta"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bbVvR8I47tFn"
      },
      "source": [
        "_, noise_key = random.split(beta_init_key)\n",
        "noise = random.normal(noise_key, (5,))*0.01"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ffbQqPlK8G25"
      },
      "source": [
        "beta_init = beta + noise"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kUep2JfO8j2N",
        "outputId": "fc48458b-2f10-484b-eeb5-810fb4014d01"
      },
      "source": [
        "beta_init"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([ 1.9979278e+00,  1.9895447e+00,  2.4958886e-02,\n",
              "              2.7045694e-03, -1.0351079e-04], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 96
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5ZJrDom68YH4"
      },
      "source": [
        "converged = False\n",
        "beta_old = beta_init\n",
        "beta_new = beta_init\n",
        "soft_threshold_partial = lambda x: soft_threshold(x, lr*penalty)\n",
        "current_iter = 0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b9kVFkXO8bIJ"
      },
      "source": [
        "beta_copy = beta_old \n",
        "current_loss = loss(beta_copy) \n",
        "current_grad = custom_grad(beta_copy)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "77QPw2T68dZw",
        "outputId": "ce11768e-2265-48db-adda-e244367314ed"
      },
      "source": [
        "current_loss + penalty * jnp.linalg.norm(beta_copy,1)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(4.2238173, dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 95
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o3SG5L5k8fUJ"
      },
      "source": [
        "w = beta_copy - lr*current_grad"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BsjookZc86N2"
      },
      "source": [
        "beta_new = jax.vmap(soft_threshold_partial, 0)(w)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZtYwKXuV88-D",
        "outputId": "f5c0e27a-2017-4f70-da21-4124ec4ced89"
      },
      "source": [
        "beta_new"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([1.9023663, 1.893549 , 0.       , 0.       , 0.       ], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 99
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qCEGLH3j8_kG"
      },
      "source": [
        "That seems right as well! Now to plug in everything and see how to goes. First we still start with the easy one."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D5cBvY7y8-H_"
      },
      "source": [
        "def proximal_methods(beta_init, max_iter, eps, lr, penalty):\n",
        "    converged = False\n",
        "    beta_old = beta_init\n",
        "    beta_new = beta_init\n",
        "    soft_threshold_partial = lambda x: soft_threshold(x, lr*penalty)\n",
        "    current_iter = 0\n",
        "    new_loss = 100000\n",
        "    while not converged and current_iter < max_iter:\n",
        "        print(\"Current iteration is %d\"% current_iter)\n",
        "        beta_copy = beta_old \n",
        "        current_loss = loss(beta_copy) + penalty*jnp.linalg.norm(beta_copy, 1)\n",
        "        current_grad = custom_grad(beta_copy)\n",
        "        w = beta_copy - lr*current_grad\n",
        "        beta_new = jax.vmap(soft_threshold_partial, 0)(w)\n",
        "        new_loss = loss(beta_new) + penalty*jnp.linalg.norm(beta_new, 1)\n",
        "        diff = jnp.abs(new_loss-current_loss)\n",
        "        print(\"The difference is %.5f\"%diff)\n",
        "        beta_old = beta_new\n",
        "        if diff <= eps:   \n",
        "            converged = True\n",
        "            print(\"Algorithm converged\")\n",
        "            break\n",
        "        else:\n",
        "            current_iter +=1\n",
        "            if current_iter >= max_iter:\n",
        "                print(\"The algorithm have failed to converge.\")\n",
        "                break\n",
        "\n",
        "    return beta_new, new_loss, converged"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VEssEqcp9jPj"
      },
      "source": [
        "max_iter = 1000\n",
        "eps = 1e-4\n",
        "lr = 0.05\n",
        "penalty = 0.1# This can be only done via trial and error"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J5ZB3Sez9o4F"
      },
      "source": [
        "key = random.PRNGKey(0)\n",
        "x_key, beta_init_key = random.split(key,2)\n",
        "x = random.normal(x_key, (10000, 5))\n",
        "beta = jnp.array([2.0,2.0,0.0,0.0,0.0])\n",
        "beta_init = beta+ random.normal(beta_init_key, (5,))*0.01\n",
        "y = (sigmoid(x.dot(beta))>=0.5).astype(jnp.float32)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NjD23yFJ-HEe",
        "outputId": "54986a4f-b27a-461e-912b-149730a85b0a"
      },
      "source": [
        "beta_sol, new_loss, converged = proximal_methods(beta_init, max_iter, eps, lr, penalty)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Current iteration is 0\n",
            "The difference is 0.00184\n",
            "Current iteration is 1\n",
            "The difference is 0.00181\n",
            "Current iteration is 2\n",
            "The difference is 0.00100\n",
            "Current iteration is 3\n",
            "The difference is 0.00083\n",
            "Current iteration is 4\n",
            "The difference is 0.00035\n",
            "Current iteration is 5\n",
            "The difference is 0.00032\n",
            "Current iteration is 6\n",
            "The difference is 0.00032\n",
            "Current iteration is 7\n",
            "The difference is 0.00032\n",
            "Current iteration is 8\n",
            "The difference is 0.00032\n",
            "Current iteration is 9\n",
            "The difference is 0.00032\n",
            "Current iteration is 10\n",
            "The difference is 0.00032\n",
            "Current iteration is 11\n",
            "The difference is 0.00032\n",
            "Current iteration is 12\n",
            "The difference is 0.00032\n",
            "Current iteration is 13\n",
            "The difference is 0.00032\n",
            "Current iteration is 14\n",
            "The difference is 0.00032\n",
            "Current iteration is 15\n",
            "The difference is 0.00031\n",
            "Current iteration is 16\n",
            "The difference is 0.00031\n",
            "Current iteration is 17\n",
            "The difference is 0.00031\n",
            "Current iteration is 18\n",
            "The difference is 0.00031\n",
            "Current iteration is 19\n",
            "The difference is 0.00031\n",
            "Current iteration is 20\n",
            "The difference is 0.00031\n",
            "Current iteration is 21\n",
            "The difference is 0.00031\n",
            "Current iteration is 22\n",
            "The difference is 0.00031\n",
            "Current iteration is 23\n",
            "The difference is 0.00031\n",
            "Current iteration is 24\n",
            "The difference is 0.00030\n",
            "Current iteration is 25\n",
            "The difference is 0.00030\n",
            "Current iteration is 26\n",
            "The difference is 0.00030\n",
            "Current iteration is 27\n",
            "The difference is 0.00030\n",
            "Current iteration is 28\n",
            "The difference is 0.00030\n",
            "Current iteration is 29\n",
            "The difference is 0.00030\n",
            "Current iteration is 30\n",
            "The difference is 0.00030\n",
            "Current iteration is 31\n",
            "The difference is 0.00030\n",
            "Current iteration is 32\n",
            "The difference is 0.00030\n",
            "Current iteration is 33\n",
            "The difference is 0.00030\n",
            "Current iteration is 34\n",
            "The difference is 0.00029\n",
            "Current iteration is 35\n",
            "The difference is 0.00029\n",
            "Current iteration is 36\n",
            "The difference is 0.00029\n",
            "Current iteration is 37\n",
            "The difference is 0.00029\n",
            "Current iteration is 38\n",
            "The difference is 0.00029\n",
            "Current iteration is 39\n",
            "The difference is 0.00029\n",
            "Current iteration is 40\n",
            "The difference is 0.00029\n",
            "Current iteration is 41\n",
            "The difference is 0.00029\n",
            "Current iteration is 42\n",
            "The difference is 0.00029\n",
            "Current iteration is 43\n",
            "The difference is 0.00028\n",
            "Current iteration is 44\n",
            "The difference is 0.00028\n",
            "Current iteration is 45\n",
            "The difference is 0.00028\n",
            "Current iteration is 46\n",
            "The difference is 0.00028\n",
            "Current iteration is 47\n",
            "The difference is 0.00028\n",
            "Current iteration is 48\n",
            "The difference is 0.00028\n",
            "Current iteration is 49\n",
            "The difference is 0.00028\n",
            "Current iteration is 50\n",
            "The difference is 0.00028\n",
            "Current iteration is 51\n",
            "The difference is 0.00028\n",
            "Current iteration is 52\n",
            "The difference is 0.00028\n",
            "Current iteration is 53\n",
            "The difference is 0.00027\n",
            "Current iteration is 54\n",
            "The difference is 0.00027\n",
            "Current iteration is 55\n",
            "The difference is 0.00027\n",
            "Current iteration is 56\n",
            "The difference is 0.00027\n",
            "Current iteration is 57\n",
            "The difference is 0.00027\n",
            "Current iteration is 58\n",
            "The difference is 0.00027\n",
            "Current iteration is 59\n",
            "The difference is 0.00027\n",
            "Current iteration is 60\n",
            "The difference is 0.00027\n",
            "Current iteration is 61\n",
            "The difference is 0.00027\n",
            "Current iteration is 62\n",
            "The difference is 0.00027\n",
            "Current iteration is 63\n",
            "The difference is 0.00026\n",
            "Current iteration is 64\n",
            "The difference is 0.00026\n",
            "Current iteration is 65\n",
            "The difference is 0.00026\n",
            "Current iteration is 66\n",
            "The difference is 0.00026\n",
            "Current iteration is 67\n",
            "The difference is 0.00026\n",
            "Current iteration is 68\n",
            "The difference is 0.00026\n",
            "Current iteration is 69\n",
            "The difference is 0.00026\n",
            "Current iteration is 70\n",
            "The difference is 0.00026\n",
            "Current iteration is 71\n",
            "The difference is 0.00026\n",
            "Current iteration is 72\n",
            "The difference is 0.00026\n",
            "Current iteration is 73\n",
            "The difference is 0.00025\n",
            "Current iteration is 74\n",
            "The difference is 0.00025\n",
            "Current iteration is 75\n",
            "The difference is 0.00025\n",
            "Current iteration is 76\n",
            "The difference is 0.00025\n",
            "Current iteration is 77\n",
            "The difference is 0.00025\n",
            "Current iteration is 78\n",
            "The difference is 0.00025\n",
            "Current iteration is 79\n",
            "The difference is 0.00025\n",
            "Current iteration is 80\n",
            "The difference is 0.00025\n",
            "Current iteration is 81\n",
            "The difference is 0.00025\n",
            "Current iteration is 82\n",
            "The difference is 0.00025\n",
            "Current iteration is 83\n",
            "The difference is 0.00024\n",
            "Current iteration is 84\n",
            "The difference is 0.00024\n",
            "Current iteration is 85\n",
            "The difference is 0.00024\n",
            "Current iteration is 86\n",
            "The difference is 0.00024\n",
            "Current iteration is 87\n",
            "The difference is 0.00024\n",
            "Current iteration is 88\n",
            "The difference is 0.00024\n",
            "Current iteration is 89\n",
            "The difference is 0.00024\n",
            "Current iteration is 90\n",
            "The difference is 0.00024\n",
            "Current iteration is 91\n",
            "The difference is 0.00024\n",
            "Current iteration is 92\n",
            "The difference is 0.00024\n",
            "Current iteration is 93\n",
            "The difference is 0.00023\n",
            "Current iteration is 94\n",
            "The difference is 0.00023\n",
            "Current iteration is 95\n",
            "The difference is 0.00023\n",
            "Current iteration is 96\n",
            "The difference is 0.00023\n",
            "Current iteration is 97\n",
            "The difference is 0.00023\n",
            "Current iteration is 98\n",
            "The difference is 0.00023\n",
            "Current iteration is 99\n",
            "The difference is 0.00023\n",
            "Current iteration is 100\n",
            "The difference is 0.00023\n",
            "Current iteration is 101\n",
            "The difference is 0.00023\n",
            "Current iteration is 102\n",
            "The difference is 0.00023\n",
            "Current iteration is 103\n",
            "The difference is 0.00022\n",
            "Current iteration is 104\n",
            "The difference is 0.00022\n",
            "Current iteration is 105\n",
            "The difference is 0.00022\n",
            "Current iteration is 106\n",
            "The difference is 0.00022\n",
            "Current iteration is 107\n",
            "The difference is 0.00022\n",
            "Current iteration is 108\n",
            "The difference is 0.00022\n",
            "Current iteration is 109\n",
            "The difference is 0.00022\n",
            "Current iteration is 110\n",
            "The difference is 0.00022\n",
            "Current iteration is 111\n",
            "The difference is 0.00022\n",
            "Current iteration is 112\n",
            "The difference is 0.00022\n",
            "Current iteration is 113\n",
            "The difference is 0.00021\n",
            "Current iteration is 114\n",
            "The difference is 0.00021\n",
            "Current iteration is 115\n",
            "The difference is 0.00021\n",
            "Current iteration is 116\n",
            "The difference is 0.00021\n",
            "Current iteration is 117\n",
            "The difference is 0.00021\n",
            "Current iteration is 118\n",
            "The difference is 0.00021\n",
            "Current iteration is 119\n",
            "The difference is 0.00021\n",
            "Current iteration is 120\n",
            "The difference is 0.00021\n",
            "Current iteration is 121\n",
            "The difference is 0.00021\n",
            "Current iteration is 122\n",
            "The difference is 0.00021\n",
            "Current iteration is 123\n",
            "The difference is 0.00021\n",
            "Current iteration is 124\n",
            "The difference is 0.00020\n",
            "Current iteration is 125\n",
            "The difference is 0.00020\n",
            "Current iteration is 126\n",
            "The difference is 0.00020\n",
            "Current iteration is 127\n",
            "The difference is 0.00020\n",
            "Current iteration is 128\n",
            "The difference is 0.00020\n",
            "Current iteration is 129\n",
            "The difference is 0.00020\n",
            "Current iteration is 130\n",
            "The difference is 0.00020\n",
            "Current iteration is 131\n",
            "The difference is 0.00020\n",
            "Current iteration is 132\n",
            "The difference is 0.00020\n",
            "Current iteration is 133\n",
            "The difference is 0.00020\n",
            "Current iteration is 134\n",
            "The difference is 0.00020\n",
            "Current iteration is 135\n",
            "The difference is 0.00019\n",
            "Current iteration is 136\n",
            "The difference is 0.00019\n",
            "Current iteration is 137\n",
            "The difference is 0.00019\n",
            "Current iteration is 138\n",
            "The difference is 0.00019\n",
            "Current iteration is 139\n",
            "The difference is 0.00019\n",
            "Current iteration is 140\n",
            "The difference is 0.00019\n",
            "Current iteration is 141\n",
            "The difference is 0.00019\n",
            "Current iteration is 142\n",
            "The difference is 0.00019\n",
            "Current iteration is 143\n",
            "The difference is 0.00019\n",
            "Current iteration is 144\n",
            "The difference is 0.00019\n",
            "Current iteration is 145\n",
            "The difference is 0.00019\n",
            "Current iteration is 146\n",
            "The difference is 0.00018\n",
            "Current iteration is 147\n",
            "The difference is 0.00018\n",
            "Current iteration is 148\n",
            "The difference is 0.00018\n",
            "Current iteration is 149\n",
            "The difference is 0.00018\n",
            "Current iteration is 150\n",
            "The difference is 0.00018\n",
            "Current iteration is 151\n",
            "The difference is 0.00018\n",
            "Current iteration is 152\n",
            "The difference is 0.00018\n",
            "Current iteration is 153\n",
            "The difference is 0.00018\n",
            "Current iteration is 154\n",
            "The difference is 0.00018\n",
            "Current iteration is 155\n",
            "The difference is 0.00018\n",
            "Current iteration is 156\n",
            "The difference is 0.00018\n",
            "Current iteration is 157\n",
            "The difference is 0.00018\n",
            "Current iteration is 158\n",
            "The difference is 0.00017\n",
            "Current iteration is 159\n",
            "The difference is 0.00017\n",
            "Current iteration is 160\n",
            "The difference is 0.00017\n",
            "Current iteration is 161\n",
            "The difference is 0.00017\n",
            "Current iteration is 162\n",
            "The difference is 0.00017\n",
            "Current iteration is 163\n",
            "The difference is 0.00017\n",
            "Current iteration is 164\n",
            "The difference is 0.00017\n",
            "Current iteration is 165\n",
            "The difference is 0.00017\n",
            "Current iteration is 166\n",
            "The difference is 0.00017\n",
            "Current iteration is 167\n",
            "The difference is 0.00017\n",
            "Current iteration is 168\n",
            "The difference is 0.00017\n",
            "Current iteration is 169\n",
            "The difference is 0.00016\n",
            "Current iteration is 170\n",
            "The difference is 0.00016\n",
            "Current iteration is 171\n",
            "The difference is 0.00016\n",
            "Current iteration is 172\n",
            "The difference is 0.00016\n",
            "Current iteration is 173\n",
            "The difference is 0.00016\n",
            "Current iteration is 174\n",
            "The difference is 0.00016\n",
            "Current iteration is 175\n",
            "The difference is 0.00016\n",
            "Current iteration is 176\n",
            "The difference is 0.00016\n",
            "Current iteration is 177\n",
            "The difference is 0.00016\n",
            "Current iteration is 178\n",
            "The difference is 0.00016\n",
            "Current iteration is 179\n",
            "The difference is 0.00016\n",
            "Current iteration is 180\n",
            "The difference is 0.00016\n",
            "Current iteration is 181\n",
            "The difference is 0.00016\n",
            "Current iteration is 182\n",
            "The difference is 0.00015\n",
            "Current iteration is 183\n",
            "The difference is 0.00015\n",
            "Current iteration is 184\n",
            "The difference is 0.00015\n",
            "Current iteration is 185\n",
            "The difference is 0.00015\n",
            "Current iteration is 186\n",
            "The difference is 0.00015\n",
            "Current iteration is 187\n",
            "The difference is 0.00015\n",
            "Current iteration is 188\n",
            "The difference is 0.00015\n",
            "Current iteration is 189\n",
            "The difference is 0.00015\n",
            "Current iteration is 190\n",
            "The difference is 0.00015\n",
            "Current iteration is 191\n",
            "The difference is 0.00015\n",
            "Current iteration is 192\n",
            "The difference is 0.00015\n",
            "Current iteration is 193\n",
            "The difference is 0.00015\n",
            "Current iteration is 194\n",
            "The difference is 0.00014\n",
            "Current iteration is 195\n",
            "The difference is 0.00014\n",
            "Current iteration is 196\n",
            "The difference is 0.00014\n",
            "Current iteration is 197\n",
            "The difference is 0.00014\n",
            "Current iteration is 198\n",
            "The difference is 0.00014\n",
            "Current iteration is 199\n",
            "The difference is 0.00014\n",
            "Current iteration is 200\n",
            "The difference is 0.00014\n",
            "Current iteration is 201\n",
            "The difference is 0.00014\n",
            "Current iteration is 202\n",
            "The difference is 0.00014\n",
            "Current iteration is 203\n",
            "The difference is 0.00014\n",
            "Current iteration is 204\n",
            "The difference is 0.00014\n",
            "Current iteration is 205\n",
            "The difference is 0.00014\n",
            "Current iteration is 206\n",
            "The difference is 0.00014\n",
            "Current iteration is 207\n",
            "The difference is 0.00014\n",
            "Current iteration is 208\n",
            "The difference is 0.00013\n",
            "Current iteration is 209\n",
            "The difference is 0.00013\n",
            "Current iteration is 210\n",
            "The difference is 0.00013\n",
            "Current iteration is 211\n",
            "The difference is 0.00013\n",
            "Current iteration is 212\n",
            "The difference is 0.00013\n",
            "Current iteration is 213\n",
            "The difference is 0.00013\n",
            "Current iteration is 214\n",
            "The difference is 0.00013\n",
            "Current iteration is 215\n",
            "The difference is 0.00013\n",
            "Current iteration is 216\n",
            "The difference is 0.00013\n",
            "Current iteration is 217\n",
            "The difference is 0.00013\n",
            "Current iteration is 218\n",
            "The difference is 0.00013\n",
            "Current iteration is 219\n",
            "The difference is 0.00013\n",
            "Current iteration is 220\n",
            "The difference is 0.00013\n",
            "Current iteration is 221\n",
            "The difference is 0.00012\n",
            "Current iteration is 222\n",
            "The difference is 0.00012\n",
            "Current iteration is 223\n",
            "The difference is 0.00012\n",
            "Current iteration is 224\n",
            "The difference is 0.00012\n",
            "Current iteration is 225\n",
            "The difference is 0.00012\n",
            "Current iteration is 226\n",
            "The difference is 0.00012\n",
            "Current iteration is 227\n",
            "The difference is 0.00012\n",
            "Current iteration is 228\n",
            "The difference is 0.00012\n",
            "Current iteration is 229\n",
            "The difference is 0.00012\n",
            "Current iteration is 230\n",
            "The difference is 0.00012\n",
            "Current iteration is 231\n",
            "The difference is 0.00012\n",
            "Current iteration is 232\n",
            "The difference is 0.00012\n",
            "Current iteration is 233\n",
            "The difference is 0.00012\n",
            "Current iteration is 234\n",
            "The difference is 0.00012\n",
            "Current iteration is 235\n",
            "The difference is 0.00012\n",
            "Current iteration is 236\n",
            "The difference is 0.00011\n",
            "Current iteration is 237\n",
            "The difference is 0.00011\n",
            "Current iteration is 238\n",
            "The difference is 0.00011\n",
            "Current iteration is 239\n",
            "The difference is 0.00011\n",
            "Current iteration is 240\n",
            "The difference is 0.00011\n",
            "Current iteration is 241\n",
            "The difference is 0.00011\n",
            "Current iteration is 242\n",
            "The difference is 0.00011\n",
            "Current iteration is 243\n",
            "The difference is 0.00011\n",
            "Current iteration is 244\n",
            "The difference is 0.00011\n",
            "Current iteration is 245\n",
            "The difference is 0.00011\n",
            "Current iteration is 246\n",
            "The difference is 0.00011\n",
            "Current iteration is 247\n",
            "The difference is 0.00011\n",
            "Current iteration is 248\n",
            "The difference is 0.00011\n",
            "Current iteration is 249\n",
            "The difference is 0.00011\n",
            "Current iteration is 250\n",
            "The difference is 0.00011\n",
            "Current iteration is 251\n",
            "The difference is 0.00010\n",
            "Current iteration is 252\n",
            "The difference is 0.00010\n",
            "Current iteration is 253\n",
            "The difference is 0.00010\n",
            "Current iteration is 254\n",
            "The difference is 0.00010\n",
            "Current iteration is 255\n",
            "The difference is 0.00010\n",
            "Current iteration is 256\n",
            "The difference is 0.00010\n",
            "Current iteration is 257\n",
            "The difference is 0.00010\n",
            "Current iteration is 258\n",
            "The difference is 0.00010\n",
            "Current iteration is 259\n",
            "The difference is 0.00010\n",
            "Algorithm converged\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aTxX90BA-Q5P",
        "outputId": "91bc1369-fc15-441f-e077-93493930f370"
      },
      "source": [
        "beta_sol"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([1.4203815, 1.3840511, 0.       , 0.       , 0.       ], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 214
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DyGaQjiCIiMs",
        "outputId": "b0618fcd-f23f-48d4-ec7c-317866514689"
      },
      "source": [
        "new_loss"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(0.55273294, dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 215
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "njYnCSfCAt7g"
      },
      "source": [
        "Now for the real test"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3f7rXQIa-o8a"
      },
      "source": [
        "key = random.PRNGKey(0)\n",
        "x_key, beta_init_key = random.split(key,2)\n",
        "x = random.normal(x_key, (10000, 5))\n",
        "beta = jnp.array([2.0,2.0,0.0,0.0,0.0])\n",
        "beta_init = random.normal(beta_init_key, (5,)) # This time beta_init is totally stochastic\n",
        "y = (sigmoid(x.dot(beta))>=0.5).astype(jnp.float32)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d-pYAS3TA4Q6"
      },
      "source": [
        "max_iter = 1000\n",
        "eps = 1e-4\n",
        "lr = 0.05\n",
        "penalty = 0.1 # This can be only done via trial and error"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sf6_rb54A9ON",
        "outputId": "93d75cbd-3aba-4fa7-f28e-89b02af35791"
      },
      "source": [
        "beta_sol, new_loss, converged = proximal_methods(beta_init, max_iter, eps, lr, penalty)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Current iteration is 0\n",
            "The difference is 0.03999\n",
            "Current iteration is 1\n",
            "The difference is 0.03988\n",
            "Current iteration is 2\n",
            "The difference is 0.03977\n",
            "Current iteration is 3\n",
            "The difference is 0.03964\n",
            "Current iteration is 4\n",
            "The difference is 0.03953\n",
            "Current iteration is 5\n",
            "The difference is 0.03940\n",
            "Current iteration is 6\n",
            "The difference is 0.03928\n",
            "Current iteration is 7\n",
            "The difference is 0.03915\n",
            "Current iteration is 8\n",
            "The difference is 0.03901\n",
            "Current iteration is 9\n",
            "The difference is 0.03888\n",
            "Current iteration is 10\n",
            "The difference is 0.03874\n",
            "Current iteration is 11\n",
            "The difference is 0.03860\n",
            "Current iteration is 12\n",
            "The difference is 0.03845\n",
            "Current iteration is 13\n",
            "The difference is 0.03830\n",
            "Current iteration is 14\n",
            "The difference is 0.03814\n",
            "Current iteration is 15\n",
            "The difference is 0.03799\n",
            "Current iteration is 16\n",
            "The difference is 0.03783\n",
            "Current iteration is 17\n",
            "The difference is 0.03767\n",
            "Current iteration is 18\n",
            "The difference is 0.03750\n",
            "Current iteration is 19\n",
            "The difference is 0.03732\n",
            "Current iteration is 20\n",
            "The difference is 0.03715\n",
            "Current iteration is 21\n",
            "The difference is 0.03697\n",
            "Current iteration is 22\n",
            "The difference is 0.03678\n",
            "Current iteration is 23\n",
            "The difference is 0.03659\n",
            "Current iteration is 24\n",
            "The difference is 0.03640\n",
            "Current iteration is 25\n",
            "The difference is 0.03620\n",
            "Current iteration is 26\n",
            "The difference is 0.03600\n",
            "Current iteration is 27\n",
            "The difference is 0.03579\n",
            "Current iteration is 28\n",
            "The difference is 0.03558\n",
            "Current iteration is 29\n",
            "The difference is 0.03536\n",
            "Current iteration is 30\n",
            "The difference is 0.03514\n",
            "Current iteration is 31\n",
            "The difference is 0.03491\n",
            "Current iteration is 32\n",
            "The difference is 0.03468\n",
            "Current iteration is 33\n",
            "The difference is 0.03445\n",
            "Current iteration is 34\n",
            "The difference is 0.03420\n",
            "Current iteration is 35\n",
            "The difference is 0.03396\n",
            "Current iteration is 36\n",
            "The difference is 0.03370\n",
            "Current iteration is 37\n",
            "The difference is 0.03344\n",
            "Current iteration is 38\n",
            "The difference is 0.03318\n",
            "Current iteration is 39\n",
            "The difference is 0.03291\n",
            "Current iteration is 40\n",
            "The difference is 0.03264\n",
            "Current iteration is 41\n",
            "The difference is 0.03236\n",
            "Current iteration is 42\n",
            "The difference is 0.03207\n",
            "Current iteration is 43\n",
            "The difference is 0.03178\n",
            "Current iteration is 44\n",
            "The difference is 0.03148\n",
            "Current iteration is 45\n",
            "The difference is 0.03118\n",
            "Current iteration is 46\n",
            "The difference is 0.03087\n",
            "Current iteration is 47\n",
            "The difference is 0.03056\n",
            "Current iteration is 48\n",
            "The difference is 0.03024\n",
            "Current iteration is 49\n",
            "The difference is 0.02992\n",
            "Current iteration is 50\n",
            "The difference is 0.02959\n",
            "Current iteration is 51\n",
            "The difference is 0.02926\n",
            "Current iteration is 52\n",
            "The difference is 0.02892\n",
            "Current iteration is 53\n",
            "The difference is 0.02858\n",
            "Current iteration is 54\n",
            "The difference is 0.02823\n",
            "Current iteration is 55\n",
            "The difference is 0.02788\n",
            "Current iteration is 56\n",
            "The difference is 0.02752\n",
            "Current iteration is 57\n",
            "The difference is 0.02716\n",
            "Current iteration is 58\n",
            "The difference is 0.02679\n",
            "Current iteration is 59\n",
            "The difference is 0.02642\n",
            "Current iteration is 60\n",
            "The difference is 0.02605\n",
            "Current iteration is 61\n",
            "The difference is 0.02172\n",
            "Current iteration is 62\n",
            "The difference is 0.01969\n",
            "Current iteration is 63\n",
            "The difference is 0.01942\n",
            "Current iteration is 64\n",
            "The difference is 0.01914\n",
            "Current iteration is 65\n",
            "The difference is 0.01887\n",
            "Current iteration is 66\n",
            "The difference is 0.01859\n",
            "Current iteration is 67\n",
            "The difference is 0.01831\n",
            "Current iteration is 68\n",
            "The difference is 0.01803\n",
            "Current iteration is 69\n",
            "The difference is 0.01775\n",
            "Current iteration is 70\n",
            "The difference is 0.01747\n",
            "Current iteration is 71\n",
            "The difference is 0.01719\n",
            "Current iteration is 72\n",
            "The difference is 0.01690\n",
            "Current iteration is 73\n",
            "The difference is 0.01662\n",
            "Current iteration is 74\n",
            "The difference is 0.01634\n",
            "Current iteration is 75\n",
            "The difference is 0.01605\n",
            "Current iteration is 76\n",
            "The difference is 0.01577\n",
            "Current iteration is 77\n",
            "The difference is 0.01549\n",
            "Current iteration is 78\n",
            "The difference is 0.01521\n",
            "Current iteration is 79\n",
            "The difference is 0.01493\n",
            "Current iteration is 80\n",
            "The difference is 0.01465\n",
            "Current iteration is 81\n",
            "The difference is 0.00985\n",
            "Current iteration is 82\n",
            "The difference is 0.00871\n",
            "Current iteration is 83\n",
            "The difference is 0.00855\n",
            "Current iteration is 84\n",
            "The difference is 0.00840\n",
            "Current iteration is 85\n",
            "The difference is 0.00824\n",
            "Current iteration is 86\n",
            "The difference is 0.00808\n",
            "Current iteration is 87\n",
            "The difference is 0.00793\n",
            "Current iteration is 88\n",
            "The difference is 0.00778\n",
            "Current iteration is 89\n",
            "The difference is 0.00762\n",
            "Current iteration is 90\n",
            "The difference is 0.00747\n",
            "Current iteration is 91\n",
            "The difference is 0.00733\n",
            "Current iteration is 92\n",
            "The difference is 0.00718\n",
            "Current iteration is 93\n",
            "The difference is 0.00703\n",
            "Current iteration is 94\n",
            "The difference is 0.00689\n",
            "Current iteration is 95\n",
            "The difference is 0.00675\n",
            "Current iteration is 96\n",
            "The difference is 0.00661\n",
            "Current iteration is 97\n",
            "The difference is 0.00647\n",
            "Current iteration is 98\n",
            "The difference is 0.00633\n",
            "Current iteration is 99\n",
            "The difference is 0.00620\n",
            "Current iteration is 100\n",
            "The difference is 0.00607\n",
            "Current iteration is 101\n",
            "The difference is 0.00594\n",
            "Current iteration is 102\n",
            "The difference is 0.00581\n",
            "Current iteration is 103\n",
            "The difference is 0.00568\n",
            "Current iteration is 104\n",
            "The difference is 0.00556\n",
            "Current iteration is 105\n",
            "The difference is 0.00544\n",
            "Current iteration is 106\n",
            "The difference is 0.00532\n",
            "Current iteration is 107\n",
            "The difference is 0.00520\n",
            "Current iteration is 108\n",
            "The difference is 0.00508\n",
            "Current iteration is 109\n",
            "The difference is 0.00497\n",
            "Current iteration is 110\n",
            "The difference is 0.00486\n",
            "Current iteration is 111\n",
            "The difference is 0.00475\n",
            "Current iteration is 112\n",
            "The difference is 0.00464\n",
            "Current iteration is 113\n",
            "The difference is 0.00454\n",
            "Current iteration is 114\n",
            "The difference is 0.00443\n",
            "Current iteration is 115\n",
            "The difference is 0.00387\n",
            "Current iteration is 116\n",
            "The difference is 0.00373\n",
            "Current iteration is 117\n",
            "The difference is 0.00365\n",
            "Current iteration is 118\n",
            "The difference is 0.00356\n",
            "Current iteration is 119\n",
            "The difference is 0.00348\n",
            "Current iteration is 120\n",
            "The difference is 0.00341\n",
            "Current iteration is 121\n",
            "The difference is 0.00333\n",
            "Current iteration is 122\n",
            "The difference is 0.00326\n",
            "Current iteration is 123\n",
            "The difference is 0.00318\n",
            "Current iteration is 124\n",
            "The difference is 0.00311\n",
            "Current iteration is 125\n",
            "The difference is 0.00304\n",
            "Current iteration is 126\n",
            "The difference is 0.00297\n",
            "Current iteration is 127\n",
            "The difference is 0.00291\n",
            "Current iteration is 128\n",
            "The difference is 0.00284\n",
            "Current iteration is 129\n",
            "The difference is 0.00278\n",
            "Current iteration is 130\n",
            "The difference is 0.00272\n",
            "Current iteration is 131\n",
            "The difference is 0.00265\n",
            "Current iteration is 132\n",
            "The difference is 0.00260\n",
            "Current iteration is 133\n",
            "The difference is 0.00221\n",
            "Current iteration is 134\n",
            "The difference is 0.00202\n",
            "Current iteration is 135\n",
            "The difference is 0.00198\n",
            "Current iteration is 136\n",
            "The difference is 0.00193\n",
            "Current iteration is 137\n",
            "The difference is 0.00189\n",
            "Current iteration is 138\n",
            "The difference is 0.00185\n",
            "Current iteration is 139\n",
            "The difference is 0.00181\n",
            "Current iteration is 140\n",
            "The difference is 0.00177\n",
            "Current iteration is 141\n",
            "The difference is 0.00173\n",
            "Current iteration is 142\n",
            "The difference is 0.00169\n",
            "Current iteration is 143\n",
            "The difference is 0.00165\n",
            "Current iteration is 144\n",
            "The difference is 0.00162\n",
            "Current iteration is 145\n",
            "The difference is 0.00158\n",
            "Current iteration is 146\n",
            "The difference is 0.00155\n",
            "Current iteration is 147\n",
            "The difference is 0.00152\n",
            "Current iteration is 148\n",
            "The difference is 0.00148\n",
            "Current iteration is 149\n",
            "The difference is 0.00145\n",
            "Current iteration is 150\n",
            "The difference is 0.00142\n",
            "Current iteration is 151\n",
            "The difference is 0.00139\n",
            "Current iteration is 152\n",
            "The difference is 0.00136\n",
            "Current iteration is 153\n",
            "The difference is 0.00133\n",
            "Current iteration is 154\n",
            "The difference is 0.00130\n",
            "Current iteration is 155\n",
            "The difference is 0.00127\n",
            "Current iteration is 156\n",
            "The difference is 0.00125\n",
            "Current iteration is 157\n",
            "The difference is 0.00122\n",
            "Current iteration is 158\n",
            "The difference is 0.00120\n",
            "Current iteration is 159\n",
            "The difference is 0.00117\n",
            "Current iteration is 160\n",
            "The difference is 0.00115\n",
            "Current iteration is 161\n",
            "The difference is 0.00112\n",
            "Current iteration is 162\n",
            "The difference is 0.00110\n",
            "Current iteration is 163\n",
            "The difference is 0.00108\n",
            "Current iteration is 164\n",
            "The difference is 0.00106\n",
            "Current iteration is 165\n",
            "The difference is 0.00103\n",
            "Current iteration is 166\n",
            "The difference is 0.00101\n",
            "Current iteration is 167\n",
            "The difference is 0.00099\n",
            "Current iteration is 168\n",
            "The difference is 0.00097\n",
            "Current iteration is 169\n",
            "The difference is 0.00095\n",
            "Current iteration is 170\n",
            "The difference is 0.00093\n",
            "Current iteration is 171\n",
            "The difference is 0.00091\n",
            "Current iteration is 172\n",
            "The difference is 0.00090\n",
            "Current iteration is 173\n",
            "The difference is 0.00088\n",
            "Current iteration is 174\n",
            "The difference is 0.00070\n",
            "Current iteration is 175\n",
            "The difference is 0.00036\n",
            "Current iteration is 176\n",
            "The difference is 0.00035\n",
            "Current iteration is 177\n",
            "The difference is 0.00035\n",
            "Current iteration is 178\n",
            "The difference is 0.00034\n",
            "Current iteration is 179\n",
            "The difference is 0.00033\n",
            "Current iteration is 180\n",
            "The difference is 0.00033\n",
            "Current iteration is 181\n",
            "The difference is 0.00032\n",
            "Current iteration is 182\n",
            "The difference is 0.00032\n",
            "Current iteration is 183\n",
            "The difference is 0.00031\n",
            "Current iteration is 184\n",
            "The difference is 0.00031\n",
            "Current iteration is 185\n",
            "The difference is 0.00030\n",
            "Current iteration is 186\n",
            "The difference is 0.00030\n",
            "Current iteration is 187\n",
            "The difference is 0.00029\n",
            "Current iteration is 188\n",
            "The difference is 0.00029\n",
            "Current iteration is 189\n",
            "The difference is 0.00028\n",
            "Current iteration is 190\n",
            "The difference is 0.00028\n",
            "Current iteration is 191\n",
            "The difference is 0.00027\n",
            "Current iteration is 192\n",
            "The difference is 0.00027\n",
            "Current iteration is 193\n",
            "The difference is 0.00026\n",
            "Current iteration is 194\n",
            "The difference is 0.00026\n",
            "Current iteration is 195\n",
            "The difference is 0.00026\n",
            "Current iteration is 196\n",
            "The difference is 0.00025\n",
            "Current iteration is 197\n",
            "The difference is 0.00025\n",
            "Current iteration is 198\n",
            "The difference is 0.00024\n",
            "Current iteration is 199\n",
            "The difference is 0.00024\n",
            "Current iteration is 200\n",
            "The difference is 0.00024\n",
            "Current iteration is 201\n",
            "The difference is 0.00023\n",
            "Current iteration is 202\n",
            "The difference is 0.00023\n",
            "Current iteration is 203\n",
            "The difference is 0.00023\n",
            "Current iteration is 204\n",
            "The difference is 0.00022\n",
            "Current iteration is 205\n",
            "The difference is 0.00022\n",
            "Current iteration is 206\n",
            "The difference is 0.00021\n",
            "Current iteration is 207\n",
            "The difference is 0.00021\n",
            "Current iteration is 208\n",
            "The difference is 0.00021\n",
            "Current iteration is 209\n",
            "The difference is 0.00020\n",
            "Current iteration is 210\n",
            "The difference is 0.00020\n",
            "Current iteration is 211\n",
            "The difference is 0.00020\n",
            "Current iteration is 212\n",
            "The difference is 0.00020\n",
            "Current iteration is 213\n",
            "The difference is 0.00019\n",
            "Current iteration is 214\n",
            "The difference is 0.00019\n",
            "Current iteration is 215\n",
            "The difference is 0.00019\n",
            "Current iteration is 216\n",
            "The difference is 0.00018\n",
            "Current iteration is 217\n",
            "The difference is 0.00018\n",
            "Current iteration is 218\n",
            "The difference is 0.00018\n",
            "Current iteration is 219\n",
            "The difference is 0.00018\n",
            "Current iteration is 220\n",
            "The difference is 0.00017\n",
            "Current iteration is 221\n",
            "The difference is 0.00017\n",
            "Current iteration is 222\n",
            "The difference is 0.00017\n",
            "Current iteration is 223\n",
            "The difference is 0.00017\n",
            "Current iteration is 224\n",
            "The difference is 0.00016\n",
            "Current iteration is 225\n",
            "The difference is 0.00016\n",
            "Current iteration is 226\n",
            "The difference is 0.00016\n",
            "Current iteration is 227\n",
            "The difference is 0.00016\n",
            "Current iteration is 228\n",
            "The difference is 0.00015\n",
            "Current iteration is 229\n",
            "The difference is 0.00015\n",
            "Current iteration is 230\n",
            "The difference is 0.00015\n",
            "Current iteration is 231\n",
            "The difference is 0.00015\n",
            "Current iteration is 232\n",
            "The difference is 0.00014\n",
            "Current iteration is 233\n",
            "The difference is 0.00014\n",
            "Current iteration is 234\n",
            "The difference is 0.00014\n",
            "Current iteration is 235\n",
            "The difference is 0.00014\n",
            "Current iteration is 236\n",
            "The difference is 0.00014\n",
            "Current iteration is 237\n",
            "The difference is 0.00013\n",
            "Current iteration is 238\n",
            "The difference is 0.00013\n",
            "Current iteration is 239\n",
            "The difference is 0.00013\n",
            "Current iteration is 240\n",
            "The difference is 0.00013\n",
            "Current iteration is 241\n",
            "The difference is 0.00013\n",
            "Current iteration is 242\n",
            "The difference is 0.00012\n",
            "Current iteration is 243\n",
            "The difference is 0.00012\n",
            "Current iteration is 244\n",
            "The difference is 0.00012\n",
            "Current iteration is 245\n",
            "The difference is 0.00012\n",
            "Current iteration is 246\n",
            "The difference is 0.00012\n",
            "Current iteration is 247\n",
            "The difference is 0.00012\n",
            "Current iteration is 248\n",
            "The difference is 0.00011\n",
            "Current iteration is 249\n",
            "The difference is 0.00011\n",
            "Current iteration is 250\n",
            "The difference is 0.00011\n",
            "Current iteration is 251\n",
            "The difference is 0.00011\n",
            "Current iteration is 252\n",
            "The difference is 0.00011\n",
            "Current iteration is 253\n",
            "The difference is 0.00011\n",
            "Current iteration is 254\n",
            "The difference is 0.00011\n",
            "Current iteration is 255\n",
            "The difference is 0.00010\n",
            "Current iteration is 256\n",
            "The difference is 0.00010\n",
            "Current iteration is 257\n",
            "The difference is 0.00010\n",
            "Current iteration is 258\n",
            "The difference is 0.00010\n",
            "Algorithm converged\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ORuRERk0A-0W",
        "outputId": "eb0e2788-5a5c-4dde-ff93-62bb66028de9"
      },
      "source": [
        "beta_sol"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([0.7655314 , 0.70897543, 0.        , 0.        , 0.        ],            dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 220
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XX8uZSuqIzVJ",
        "outputId": "03a01646-fdc9-4c53-8389-4b05555a3acd"
      },
      "source": [
        "new_loss"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(0.54692423, dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 221
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2RDN3Ct7BSA9"
      },
      "source": [
        "As you can see, it is usually not the case that you will get any good results considering that everything does not have to work so well in the numerical sense. As you can see, if we start from the real value we may end up somewhere else. In fact, it is perfectly reasonable to assume in the first situation, the algorithm has not converged at all, since a better solution is found in the second instance. \n",
        "\n",
        "**Exercise: 1** There is an improved version of line search in lasso.pdf. Try to implement it and see if it helps.\n",
        "\n",
        "**Exercise: 2** One way to alleviate the problem is by giving a better init function. Usually it is a little bit cheating to assume we have already a logistic regression function, so try a linear regression function and see what happens. \n",
        "\n",
        "**Exercise: 3 (Difficult)** Try to print out everything and see what might go wrong in the algorithm and see one can improve there. Note that without the soft threshold, what we have just implemented is basically a simple gradient decent which can be pretty bad. Is there anyway to improve it?\n",
        "\n"
      ]
    }
  ]
}